package ai.knowly.langtorch.store.vectordb.integration.pgvector;

import ai.knowly.langtorch.store.vectordb.integration.pgvector.exception.PGVectorServiceInstantiationException;
import com.google.common.flogger.FluentLogger;
import com.pgvector.PGvector;

import javax.inject.Inject;
import java.sql.*;

/**
 * The PGVectorService class provides methods for interacting with the PostgreSQL Vector extension.
 * It allows executing SQL statements, preparing statements, and querying the database.
 */
public class PGVectorService {

  private static final String CREATE_VECTOR_EXTENSION_QUERY =
      "CREATE EXTENSION IF NOT EXISTS vector";
  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
  private final Connection connection;
  private final Statement defaultStatement;

  /**
   * Constructs a new PGVectorService instance with the provided database connection.
   *
   * @param connection the database connection
   */
  @Inject
  public PGVectorService(Connection connection) {
    this.connection = connection;
    try {
      PGvector.addVectorType(connection);
      defaultStatement = connection.createStatement();
      defaultStatement.executeUpdate(CREATE_VECTOR_EXTENSION_QUERY);
    } catch (SQLException e) {
      logger.atSevere().withCause(e).log("Error while initialising PGVectorService");
      throw new PGVectorServiceInstantiationException(e);
    }
  }

  /**
   * Executes the given SQL statement and returns the number of affected rows.
   *
   * @param sql the SQL statement to execute
   * @return the number of affected rows or 0 for SQL statements that return nothing
   * @throws SQLException if a database access error occurs or the SQL statement is invalid
   */
  public int executeUpdate(String sql) throws SQLException {
    return defaultStatement.executeUpdate(sql);
  }

  /**
   * Creates a PreparedStatement object for sending parameterized SQL statements to the database.
   *
   * @param sql the SQL statement to prepare
   * @return a new PreparedStatement object
   * @throws SQLException if a database access error occurs or the SQL statement is invalid
   */
  public PreparedStatement prepareStatement(String sql) throws SQLException {
    return connection.prepareStatement(sql);
  }

  /**
   * Executes the given SQL query and returns the ResultSet object generated by the query.
   *
   * @param sql the SQL query to execute
   * @return the ResultSet object generated by the query
   * @throws SQLException if a database access error occurs or the SQL statement is invalid
   */
  public ResultSet query(String sql) throws SQLException {
    return defaultStatement.executeQuery(sql);
  }
}
